import os
from tqdm import tqdm

import numpy as np

current_path = os.path.split(os.path.abspath(__file__))[0]
# Load the images


val_index = int((3837+1021)*0.8)

index = 0
with open('train.txt', 'w') as train_txt:
    with open('val.txt', 'w') as val_txt:
        for dir_index in sorted(os.listdir('dataset_all')):
            for image_path in tqdm(sorted(os.listdir(os.path.join('dataset_all',dir_index,'images')))):
                index  += 1
                # 如果大于某个序号，就都保存在验证集的txt中
                if index  <= val_index:
                    train_txt.write(os.path.join(current_path,'dataset_all',dir_index, 'images', image_path) + '\n')
                else:
                    val_txt.write(os.path.join(current_path,'dataset_all',dir_index, 'images', image_path) + '\n')



